from torchvision import datasets, transforms
from base import BaseDataLoader
from data_loader.dataset import Dataset2D, Dataset3D
import os
from pathlib import Path

class Data3DLoader(BaseDataLoader):
    def __init__(self, data_dir, batch_size, shuffle=False, validation_split=0.0, num_workers=6):
        self.data_dir = str(Path(os.getcwd()).parent) + data_dir
        self.dataset = Dataset3D(self.data_dir)
        super().__init__(self.dataset, batch_size, shuffle, validation_split, num_workers)

class Data2DLoader(BaseDataLoader):
    def __init__(self, data_dir, batch_size, eof_name, shuffle=False, validation_split=0.0, num_workers=6):
        self.data_dir = str(Path(os.getcwd()).parent) + data_dir
        self.dataset = Dataset2D(self.data_dir, eof_name)
        super().__init__(self.dataset, batch_size, shuffle, validation_split, num_workers)

# import torch.nn.functional as F
# import scipy.ndimage

# class ResizeTensor(object):
#     def __init__(self, output_size):
#         assert isinstance(output_size, (int, tuple))
#         self.output_size = output_size
#     def __call__(self, sample):
#         # out = F.interpolate(sample, size=self.output_size)
#         out = scipy.ndimage.interpolation.zoom(sample, [1, self.output_size/64, self.output_size/64], order=0)
#         return out

# class ResU2dLoader(BaseDataLoader):
#     """ Blood data loading"""
#     def __init__(self, data_dir, batch_size, shuffle=True, validation_split=0.0, num_workers=1, training=True):
#         # load dataset
#         self.data_dir = str(Path(os.getcwd()).parent) + data_dir
#         self.dataset = Bloodset(path=self.data_dir, train=training, transforms=ResizeTensor(256))
#         super().__init__(self.dataset, batch_size, shuffle, validation_split, num_workers)
